import argparse, os
import copy
import json
import os
import pathlib
import shutil
import time
import uuid

import librosa
import gradio as gr
import numpy as np
import soundfile

from common_utils import *

from common_utils import _root_dir

g_json_key_text = "text"
g_json_key_prompt = "prompt"
g_json_key_audio = "audio"
g_json_key_image = "image"
# g_load_file = ""
g_load_format = "resource"

g_max_json_index = 100
# g_index = 0
g_batch = 100


# g_text_list = []
# g_prompt_list = []
# g_audio_list = []
# g_image_list = []
# g_resource_list = []
# g_checkbox_list = []
# g_data_json = []


def reload_data(index, batch, g_data_json=()):
    # g_index
    g_index = index
    # g_batch
    g_batch = batch
    datas = g_data_json[index:index + batch]
    output = []
    for d in datas:
        output.append(d)
    return output


def b_change_index(index, batch, g_data_json=()):
    # g_index, g_batch
    g_index, g_batch = index, batch
    datas = reload_data(index, batch, g_data_json=g_data_json)
    output = []
    # text
    for i, _ in enumerate(datas):
        output.append(
            {
                "__type__": "update",
                "label": f"【{i + index}】文本",
                "value": _[g_json_key_text]
            }
        )
    for _ in range(g_batch - len(datas)):
        output.append(
            {
                "__type__": "update",
                "label": f"文本",
                "value": ""
            }
        )

    # prompt
    for i, _ in enumerate(datas):
        output.append(
            {
                "__type__": "update",
                "label": f"图像提示词",
                "value": _[g_json_key_prompt]
            }
        )
    for _ in range(g_batch - len(datas)):
        output.append(
            {
                "__type__": "update",
                "label": f"图像提示词",
                "value": ""
            }
        )

    # audio
    for _ in datas:
        output.append(_[g_json_key_audio])
    for _ in range(g_batch - len(datas)):
        output.append(None)

    # image
    for _ in datas:
        output.append(_[g_json_key_image])
    for _ in range(g_batch - len(datas)):
        output.append(None)

    # resource
    for _ in datas:
        output.append(_)
    for _ in range(g_batch - len(datas)):
        output.append(None)

    # check
    for _ in datas:
        output.append(True)
    for _ in range(g_batch - len(datas)):
        output.append(False)

    return output


def b_next_index(index, batch):
    # b_save_file()
    if (index + batch) <= g_max_json_index:
        return index + batch, *b_change_index(index + batch, batch)
    else:
        return index, *b_change_index(index, batch)


def b_previous_index(index, batch):
    # b_save_file()
    if (index - batch) >= 0:
        return index - batch, *b_change_index(index - batch, batch)
    else:
        return 0, *b_change_index(0, batch)


def b_submit_change(*text_list, g_data_json=(), g_index=0):
    # global g_data_json
    change = False
    for i, new_text in enumerate(text_list):
        if g_index + i <= g_max_json_index:
            new_text = new_text.strip() + ' '
            if (g_data_json[g_index + i][g_json_key_text] != new_text):
                g_data_json[g_index + i][g_json_key_text] = new_text
                change = True
    if change:
        b_save_file()
    return g_index, *b_change_index(g_index, g_batch)


def b_delete_audio(*checkbox_list, g_data_json=(), g_index=0):
    # global g_data_json, g_index, g_max_json_index
    b_save_file()
    change = False
    for i, checkbox in reversed(list(enumerate(checkbox_list))):
        if g_index + i < len(g_data_json):
            if (checkbox == True):
                g_data_json.pop(g_index + i)
                change = True

    g_max_json_index = len(g_data_json) - 1
    if g_index > g_max_json_index:
        g_index = g_max_json_index
        g_index = g_index if g_index >= 0 else 0
    if change:
        b_save_file()
    # return gr.Slider(value=g_index, maximum=(g_max_json_index if g_max_json_index>=0 else 0)), *b_change_index(g_index, g_batch)
    return {"value": g_index, "__type__": "update",
            "maximum": (g_max_json_index if g_max_json_index >= 0 else 0)}, *b_change_index(g_index, g_batch)


def b_invert_selection(*checkbox_list):
    new_list = [not item if item is True else True for item in checkbox_list]
    return new_list


def get_next_path(filename):
    base_dir = os.path.dirname(filename)
    base_name = os.path.splitext(os.path.basename(filename))[0]
    for i in range(100):
        new_path = os.path.join(base_dir, f"{base_name}_{str(i).zfill(2)}.wav")
        if not os.path.exists(new_path):
            return new_path
    return os.path.join(base_dir, f'{str(uuid.uuid4())}.wav')


def b_audio_split(audio_breakpoint, *checkbox_list, g_data_json=(), g_index):
    # global g_data_json, g_max_json_index
    checked_index = []
    for i, checkbox in enumerate(checkbox_list):
        if (checkbox == True and g_index + i < len(g_data_json)):
            checked_index.append(g_index + i)
    if len(checked_index) == 1:
        index = checked_index[0]
        audio_json = copy.deepcopy(g_data_json[index])
        path = audio_json[g_json_key_audio]
        data, sample_rate = librosa.load(path, sr=None, mono=True)
        audio_maxframe = len(data)
        break_frame = int(audio_breakpoint * sample_rate)

        if (break_frame >= 1 and break_frame < audio_maxframe):
            audio_first = data[0:break_frame]
            audio_second = data[break_frame:]
            nextpath = get_next_path(path)
            soundfile.write(nextpath, audio_second, sample_rate)
            soundfile.write(path, audio_first, sample_rate)
            g_data_json.insert(index + 1, audio_json)
            g_data_json[index + 1][g_json_key_audio] = nextpath
            b_save_file()

    # g_max_json_index = len(g_data_json) - 1
    # return gr.Slider(value=g_index, maximum=g_max_json_index), *b_change_index(g_index, g_batch)
    return {"value": g_index, "maximum": g_max_json_index, "__type__": "update"}, *b_change_index(g_index, g_batch)


def b_merge_audio(interval_r, *checkbox_list, g_data_json=(), g_index=0):
    # g_data_json, g_max_json_index
    b_save_file()
    checked_index = []
    audios_audio = []
    audios_text = []
    audios_image = []
    for i, checkbox in enumerate(checkbox_list):
        if (checkbox == True and g_index + i < len(g_data_json)):
            checked_index.append(g_index + i)

    if (len(checked_index) > 1):
        for i in checked_index:
            audios_audio.append(g_data_json[i][g_json_key_audio])
            audios_text.append(g_data_json[i][g_json_key_text])
            audios_image.append(g_data_json[i][g_json_key_image])
        for i in reversed(checked_index[1:]):
            g_data_json.pop(i)

        base_index = checked_index[0]
        base_path = audios_audio[0]
        g_data_json[base_index][g_json_key_text] = "".join(audios_text)

        audio_list = []
        l_sample_rate = None
        for i, path in enumerate(audios_audio):
            data, sample_rate = librosa.load(path, sr=l_sample_rate, mono=True)
            l_sample_rate = sample_rate
            if (i > 0):
                silence = np.zeros(int(l_sample_rate * interval_r))
                audio_list.append(silence)

            audio_list.append(data)

        audio_concat = np.concatenate(audio_list)

        soundfile.write(base_path, audio_concat, l_sample_rate)

        b_save_file()

    # g_max_json_index = len(g_data_json) - 1

    # return gr.Slider(value=g_index, maximum=g_max_json_index), *b_change_index(g_index, g_batch)
    return {"value": g_index, "maximum": g_max_json_index, "__type__": "update"}, *b_change_index(g_index, g_batch)


def b_save_json(g_data_json=(), g_load_file=''):
    with open(g_load_file, 'w', encoding="utf-8") as file:
        for data in g_data_json:
            file.write(f'{json.dumps(data, ensure_ascii=False)}\n')


def b_save_list(g_data_json=(), g_load_file=''):
    with open(g_load_file, 'w', encoding="utf-8") as file:
        for data in g_data_json:
            wav_path = data["wav_path"]
            speaker_name = data["speaker_name"]
            language = data["language"]
            text = data["text"]
            file.write(f"{wav_path}|{speaker_name}|{language}|{text}".strip() + '\n')


def b_load_json(g_load_file=''):
    # g_data_json, g_max_json_index
    with open(g_load_file, 'r', encoding="utf-8") as file:
        g_data_json = file.readlines()
        g_data_json = [json.loads(line) for line in g_data_json]
        g_max_json_index = len(g_data_json) - 1
    return g_data_json


def b_load_resource(g_load_file: object = '') -> object:
    # g_data_json, g_max_json_index
    code_name = '/'.join(pathlib.Path(g_load_file).parts[-3:-1])
    # code_name = pathlib.Path(g_load_file).parent.name
    g_data_json = []
    for res_path in sorted(pathlib.Path(g_load_file).glob('*.json'), key=lambda x: int(x.stem.split('_')[-1])):
        dt = json.load(open(res_path, encoding='utf8'))
        dt['image'] = get_abspath(code_name, dt['image'])
        dt['audio'] = get_abspath(code_name, dt['audio'])
        dt['resource'] = get_abspath(code_name, dt['resource'])
        g_data_json.append(dt)
    # g_max_json_index = len(g_data_json) - 1
    return g_data_json


def b_load_list(g_data_json, g_load_file=''):
    # g_data_json, g_max_json_index
    with open(g_load_file, 'r', encoding="utf-8") as source:
        data_list = source.readlines()
        for _ in data_list:
            data = _.split('|')
            if (len(data) == 4):
                wav_path, speaker_name, language, text = data
                g_data_json.append(
                    {
                        'wav_path': wav_path,
                        'speaker_name': speaker_name,
                        'language': language,
                        'text': text.strip()
                    }
                )
            else:
                print("error line:", data)
        g_max_json_index = len(g_data_json) - 1
    return g_data_json


def b_save_file():
    if g_load_format == "json":
        b_save_json()
    elif g_load_format == "list":
        b_save_list()


def b_load_file():
    if g_load_format == "json":
        b_load_json()
    elif g_load_format == "list":
        b_load_list()
    elif g_load_format == "resource":
        b_load_resource()


def set_global(load_json, load_list, load_resource, json_key_text, json_key_audio, json_key_image, batch):
    # g_json_key_text, g_json_key_audio, g_json_key_image, g_load_file, g_load_format, g_batch

    g_batch = int(batch)

    if (load_json != "None"):
        g_load_format = "json"
        g_load_file = load_json
    elif (load_list != "None"):
        g_load_format = "list"
        g_load_file = load_list
    else:
        g_load_format = "resource"
        g_load_file = load_resource

    g_json_key_text = json_key_text
    g_json_key_audio = json_key_audio
    g_json_key_image = json_key_image

    b_load_file()



def b_text_change(text, resource):
    if resource:
        flag = text == resource['text']
    else:
        flag = False
    return flag, flag


def b_prompt_change(prompt, resource):
    if resource:
        flag = prompt == resource['prompt']
    else:
        flag = False
    return flag


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process some integers.')
    parser.add_argument('--load_json',
                        default=r"None",
                        help='source file, like demo.json')
    parser.add_argument('--is_share', default="False", help='whether webui is_share=True')
    parser.add_argument('--load_list', default="None", help='source file, like demo.list')
    parser.add_argument('--load_resource',
                        default=r"C:\Users\kuang\github\kuangdd2024\auto-video-generateor\mnt\materials\2024-08-16_10.22.51\resource",
                        help='source file, like /resource')
    parser.add_argument('--webui_port_subfix', default=9871, help='source file, like demo.list')
    parser.add_argument('--json_key_text', default="text", help='the text key name in json, Default: text')
    parser.add_argument('--json_key_audio', default="audio", help='the path key name in json, Default: audio')
    parser.add_argument('--json_key_image', default="image", help='the path key name in json, Default: image')
    parser.add_argument('--g_batch', default=10, help='max number g_batch wav to display, Default: 10')

    args = parser.parse_args()

    set_global(args.load_json, args.load_list, args.load_resource, args.json_key_text, args.json_key_audio,
               args.json_key_image,
               args.g_batch)
